import matplotlib
matplotlib.use('TkAgg')
import pandas as pd
import numpy as np
import joblib
import matplotlib.pyplot as plt
from matplotlib import rcParams
rcParams['font.sans-serif'] = ['SimHei']   # 中文字体
rcParams['axes.unicode_minus'] = False     # 正常显示负号
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    classification_report, accuracy_score,
    f1_score, precision_score, recall_score,
    confusion_matrix, roc_auc_score, roc_curve
)
from word2vec_processor import Word2VecProcessor
from utils.config import Config

# ================= 配置路径 =================
root_path = "D:/pycode/group4_nlp_project"
conf = Config(root_path)
data_path = conf.train_path
model_path = conf.basic_model_path

label_lr_model_path = model_path + "label_lr.pkl"
label_w2v_path = model_path + "label_word2vec.pkl"


def train_label_model(data_path=data_path):
    # 1. 读取数据
    df = pd.read_csv(data_path)

    # 2. 训练 Word2Vec
    print("训练 Word2Vec 模型...")
    w2v = Word2VecProcessor(vector_size=100, window=5, min_count=2, epochs=10)
    w2v.train(df, text_column="review_clean")
    w2v.save(label_w2v_path)

    # 3. 转换句子向量
    print("生成句子向量...")
    df["vector"] = df["review_clean"].apply(lambda x: w2v.get_sentence_vector(str(x)))
    df = df[df["vector"].notnull()].reset_index(drop=True)
    X = np.vstack(df["vector"].values)
    y_label = df["label"]   # 二分类标签

    # 4. 切分数据
    X_train, X_test, y_train, y_test = train_test_split(
        X, y_label, test_size=0.2, random_state=42, stratify=y_label
    )

    # 5. 训练逻辑回归
    print("训练逻辑回归分类器...")
    lr = LogisticRegression(max_iter=200, solver="liblinear", random_state=917)
    lr.fit(X_train, y_train)

    # 6. 预测与评估
    y_pred = lr.predict(X_test)
    y_prob = lr.predict_proba(X_test)[:, 1]

    print("\n==== 二分类任务 (逻辑回归 - label) ====")
    print(classification_report(y_test, y_pred))

    # 更多指标
    acc = accuracy_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred)
    prec = precision_score(y_test, y_pred)
    rec = recall_score(y_test, y_pred)
    auc = roc_auc_score(y_test, y_prob)

    print("准确率:", acc)
    print("F1-score:", f1)
    print("Precision:", prec)
    print("Recall:", rec)
    print("AUC:", auc)

    # === 输出指标表格 ===
    metrics_df = pd.DataFrame({
        "Accuracy": [acc],
        "F1": [f1],
        "Precision": [prec],
        "Recall": [rec],
        "AUC": [auc]
    })
    print("\n==== 指标汇总 ====")
    print(metrics_df.to_string(index=False))

    # === 混淆矩阵 ===
    cm = confusion_matrix(y_test, y_pred)
    plt.figure(figsize=(6, 5))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Greens",
                xticklabels=lr.classes_, yticklabels=lr.classes_)
    plt.xlabel("预测值")
    plt.ylabel("真实值")
    plt.title("混淆矩阵 (逻辑回归 - label)")
    plt.tight_layout()
    plt.show()

    # === ROC曲线 ===
    fpr, tpr, _ = roc_curve(y_test, y_prob)
    plt.figure(figsize=(6, 5))
    plt.plot(fpr, tpr, label=f"AUC = {auc:.3f}")
    plt.plot([0, 1], [0, 1], "--", color="gray")
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title("ROC曲线 (逻辑回归 - label)")
    plt.legend()
    plt.tight_layout()
    plt.show()

    # 7. 保存模型和 Word2Vec
    joblib.dump(lr, label_lr_model_path)
    print(f"\n模型已保存: {label_lr_model_path}")
    print(f"词向量已保存: {label_w2v_path}")


if __name__ == "__main__":
    train_label_model()
